import ryu.app.simple_switch_13
from ryu.controller import ofp_event
from ryu.controller.handler import MAIN_DISPATCHER, CONFIG_DISPATCHER, set_ev_cls
from ryu.ofproto import ofproto_v1_3
from ryu.lib.packet import packet
from ryu.lib.packet import ethernet
from ryu.lib.packet import ipv4
from ryu.lib.packet import tcp
from ryu.lib.packet import ether_types
from ryu.ofproto import ether
import os
import sys
import threading
from tqdm import tqdm

CLIENT = {'ip': '10.0.1.5', 'mac': '00:00:00:00:00:03'}
SERVER1 = {'ip': '10.0.1.2', 'mac': '00:00:00:00:00:01'}
SERVER2 = {'ip': '10.0.1.3', 'mac': '00:00:00:00:00:02'}

# Global variable to record traffic analysis
traffic_analysis_log = []

class TrafficRedirectionApp(ryu.app.simple_switch_13.SimpleSwitch13):
    def __init__(self, *args, **kwargs):
        super(TrafficRedirectionApp, self).__init__(*args, **kwargs)
        self.dpid_to_port = {}
        print("TrafficRedirectionApp initialized")
        self.print_system_info()
        self.lock = threading.Lock()

    def print_system_info(self):
        """
        Print system information.
        """
        print(f"Operating System: {os.name}")
        print(f"Python Version: {sys.version}")
        print(f"Platform: {sys.platform}")

    @set_ev_cls(ofp_event.EventOFPSwitchFeatures, CONFIG_DISPATCHER)
    def switch_features_handler(self, event):
        """
        Handler for switch features event.

        Args:
            event (ryu.controller.ofp_event.EventOFPSwitchFeatures): The switch features event.
        """
        try:
            datapath = event.msg.datapath
            self.install_table_miss_flow(datapath)
        except Exception as e:
            print(f"Exception in switch_features_handler: {e}")

    def install_table_miss_flow(self, datapath):
        """
        Install table-miss flow entry.

        Args:
            datapath (ryu.controller.controller.Datapath): The datapath to install the flow entry.
        """
        try:
            proto = datapath.ofproto
            proto_parser = datapath.ofproto_parser
            flow_match = proto_parser.OFPMatch()
            flow_actions = [proto_parser.OFPActionOutput(proto.OFPP_CONTROLLER, proto.OFPCML_NO_BUFFER)]
            self.install_flow(datapath, 0, flow_match, flow_actions)
        except Exception as e:
            print(f"Exception in install_table_miss_flow: {e}")

    def install_flow(self, datapath, priority, flow_match, flow_actions, buffer_id=None, idle_timeout=5):
        """
        Add a flow entry to the datapath.

        Args:
            datapath (ryu.controller.controller.Datapath): The datapath to install the flow entry.
            priority (int): The priority of the flow entry.
            flow_match (ryu.ofproto.ofproto_v1_3_parser.OFPMatch): The match criteria for the flow entry.
            flow_actions (list): The actions to apply to the flow entry.
            buffer_id (int, optional): The buffer ID. Defaults to None.
            idle_timeout (int, optional): The idle timeout for the flow entry. Defaults to 5.
        """
        try:
            proto = datapath.ofproto
            proto_parser = datapath.ofproto_parser
            instructions = [proto_parser.OFPInstructionActions(proto.OFPIT_APPLY_ACTIONS, flow_actions)]
            flow_mod = proto_parser.OFPFlowMod(datapath=datapath, buffer_id=buffer_id, priority=priority,
                                               match=flow_match, instructions=instructions, idle_timeout=idle_timeout)
            datapath.send_msg(flow_mod)
        except Exception as e:
            print(f"Exception in install_flow: {e}")

    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)
    def packet_in_handler(self, event):
        """
        Handler for packet in event.

        Args:
            event (ryu.controller.ofp_event.EventOFPPacketIn): The packet in event.
        """
        try:
            datapath = event.msg.datapath
            proto = datapath.ofproto
            proto_parser = datapath.ofproto_parser
            in_port = event.msg.match['in_port']

            pkt = packet.Packet(event.msg.data)
            eth_pkt = pkt.get_protocols(ethernet.ethernet)[0]

            if eth_pkt.ethertype == ether_types.ETH_TYPE_LLDP:
                return

            dst = eth_pkt.dst
            src = eth_pkt.src
            dpid = datapath.id

            self.dpid_to_port.setdefault(dpid, {})
            print(f"Packet in {dpid} {src} {dst} {in_port}")
            self.dpid_to_port[dpid][src] = in_port

            out_port = self.dpid_to_port[dpid].get(dst, proto.OFPP_FLOOD)
            flow_actions = [proto_parser.OFPActionOutput(out_port)]

            if out_port != proto.OFPP_FLOOD:
                flow_match = proto_parser.OFPMatch(in_port=in_port, eth_dst=dst, eth_src=src)
                self.install_flow(datapath, 1, flow_match, flow_actions)

            data = event.msg.data if event.msg.buffer_id == proto.OFP_NO_BUFFER else None
            self.send_packet_out(datapath, event.msg.buffer_id, in_port, flow_actions, data)

            ip_pkt = pkt.get_protocol(ipv4.ipv4)
            tcp_pkt = pkt.get_protocol(tcp.tcp)
            if ip_pkt and tcp_pkt and tcp_pkt.bits == tcp.TCP_SYN:
                print("TCP SYN packet detected")
                self.redirect_traffic(datapath, in_port, ip_pkt, tcp_pkt)
        except Exception as e:
            print(f"Exception in packet_in_handler: {e}")

    def send_packet_out(self, datapath, buffer_id, in_port, flow_actions, data):
        """
        Send packet out.

        Args:
            datapath (ryu.controller.controller.Datapath): The datapath to send the packet out.
            buffer_id (int): The buffer ID.
            in_port (int): The input port.
            flow_actions (list): The actions to apply to the packet.
            data (bytes, optional): The packet data. Defaults to None.
        """
        try:
            proto = datapath.ofproto
            proto_parser = datapath.ofproto_parser
            packet_out = proto_parser.OFPPacketOut(datapath=datapath, buffer_id=buffer_id,
                                                   in_port=in_port, actions=flow_actions, data=data)
            datapath.send_msg(packet_out)
            print(f"Packet out sent from port {in_port}")
        except Exception as e:
            print(f"Exception in send_packet_out: {e}")

    def redirect_traffic(self, datapath, in_port, ip_pkt=None, tcp_pkt=None):
        """
        Redirect traffic to SERVER2.

        Args:
            datapath (ryu.controller.controller.Datapath): The datapath to redirect the traffic.
            in_port (int): The input port.
            ip_pkt (ryu.lib.packet.ipv4.ipv4, optional): The IP packet. Defaults to None.
            tcp_pkt (ryu.lib.packet.tcp.tcp, optional): The TCP packet. Defaults to None.
        """
        try:
            proto_parser = datapath.ofproto_parser
            proto = datapath.ofproto

            # Extract match fields
            eth_type = ether_types.ETH_TYPE_IP
            ipv4_src = ip_pkt.src if ip_pkt else CLIENT['ip']
            ipv4_dst = ip_pkt.dst if ip_pkt else SERVER1['ip']
            ip_proto = ip_pkt.proto if ip_pkt else 6  # Default to TCP
            tcp_src = tcp_pkt.src_port if tcp_pkt else 0
            tcp_dst = tcp_pkt.dst_port if tcp_pkt else 0

            # Create match criteria
            flow_match = proto_parser.OFPMatch(
                eth_type=eth_type,
                ipv4_src=ipv4_src,
                ipv4_dst=ipv4_dst,
                ip_proto=ip_proto,
                tcp_src=tcp_src,
                tcp_dst=tcp_dst
            )

            # Create actions
            flow_actions = [
                proto_parser.OFPActionSetField(eth_dst=SERVER2['mac']),
                proto_parser.OFPActionSetField(ipv4_dst=SERVER2['ip']),
                proto_parser.OFPActionOutput(proto.OFPP_IN_PORT)
            ]

            self.install_flow(datapath, 2, flow_match, flow_actions)

            self.send_packet_out(datapath, proto.OFP_NO_BUFFER, in_port, flow_actions, None)
            print(f"Traffic redirected to {SERVER2['ip']}")

            # Start traffic analysis in a separate thread
            analysis_thread = threading.Thread(target=self.analyze_traffic, args=(datapath, in_port, ip_pkt, tcp_pkt))
            analysis_thread.start()
        except Exception as e:
            print(f"Exception in redirect_traffic: {e}")

    def analyze_traffic(self, datapath, in_port, ip_pkt, tcp_pkt):
        """
        Analyze traffic using tqdm and record the analysis.

        Args:
            datapath (ryu.controller.controller.Datapath): The datapath to analyze the traffic.
            in_port (int): The input port.
            ip_pkt (ryu.lib.packet.ipv4.ipv4): The IP packet.
            tcp_pkt (ryu.lib.packet.tcp.tcp): The TCP packet.
        """
        try:
            print("Analyzing traffic...")
            for i in tqdm(range(100)):
                pass
            with self.lock:
                traffic_analysis_log.append(f"Analyzed traffic from {ip_pkt.src} to {ip_pkt.dst} on port {in_port}")
            print("Traffic analysis complete")
        except Exception as e:
            print(f"Exception in analyze_traffic: {e}")

    def print_traffic_analysis_log(self):
        """
        Print the traffic analysis log.
        """
        with self.lock:
            print("Traffic Analysis Log:")
            for log in traffic_analysis_log:
                print(log)

if __name__ == "__main__":
    print("Starting TrafficRedirectionApp")
    app = TrafficRedirectionApp()
    # Start a separate thread to periodically print the traffic analysis log
    log_thread = threading.Thread(target=app.print_traffic_analysis_log)
    log_thread.start()
